import asyncio
import datetime
import logging
import multiprocessing
import os
import sys
import threading
import time

import dns.resolver
from dns import message
from dns.rdataclass import RdataClass
from dns.rdatatype import RdataType
from pymongo import MongoClient

from transform import timestampToTime
fmt = '[%(asctime)s]{%(filename)s:%(lineno)d}[%(levelname)s] %(message)s'
logging.basicConfig(level=logging.DEBUG, format=fmt)
logger = logging.getLogger(__name__)

# 追加写入文件a ，设置utf-8编码防止中文写入乱码
date_str = datetime.datetime.now().strftime("%Y-%m-%d")
file_handler = logging.FileHandler(f'{date_str}.txt', 'a', encoding='utf-8')
# 向文件输出的日志级别
file_handler.setLevel(logging.DEBUG)
file_handler.setFormatter(logging.Formatter(fmt))
logger.addHandler(file_handler)

base_nameservers = ['8.8.8.8']
dns.resolver.default_resolver = dns.resolver.Resolver(configure=False)
dns.resolver.default_resolver.nameservers = base_nameservers
client = MongoClient('mongodb://lxt:aa347347@173.82.107.218:27017/?authSource=admin')
db = client['dns']
collection = db['resolvers']
collection_timeout = db['resolvers_info']

collection_domains = db['domains']

total_ip = 0
MAX_PORTS = 20000
MAX_INTERVAL = 1000
sem = multiprocessing.Semaphore(MAX_PORTS)


def getTTL(domain: str):
    resolver = dns.resolver.Resolver(configure=False)
    resolver.nameservers = base_nameservers

    name = dns.name.from_text(domain)
    logger.info(f'query name: {name}')
    res = resolver.resolve(name, 'A')

    # Parse cname chain
    while True:
        cnameSet = res.response.get_rrset(message.ANSWER, name, rdclass=RdataClass.IN, rdtype=RdataType.CNAME)
        if not cnameSet:
            break
        cname: dns.rdtypes.ANY.CNAME.CNAME = cnameSet.pop()
        name = dns.name.from_text(cname.to_text())

    parent = name.parent()  # find the NS of its parent zone
    # Find authority NS, from bottom to top, iteratively
    while True:
        try:
            res = resolver.resolve(parent, rdtype=RdataType.NS)
            break
        except dns.resolver.NXDOMAIN:
            parent = parent.parent()
        except dns.resolver.NoAnswer:
            parent = parent.parent()

    nsSet = res.response.get_rrset(message.ANSWER, parent, RdataClass.IN, RdataType.NS)
    ns: dns.rdtypes.ANY.NS.NS = nsSet.pop()  # randomly select a NS server

    # query the IP of the NS
    nsName = dns.name.from_text(ns.to_text())
    res = resolver.resolve(nsName)
    nsSet = res.response.get_rrset(message.ANSWER, nsName, RdataClass.IN, RdataType.A)
    # print(nsSet)

    # query the authority NS for the TTL of the A record
    resolver.nameservers = [ns.to_text() for ns in nsSet]
    res = resolver.resolve(name, raise_on_no_answer=False)
    if not res.response.answer:
        assert res.response.authority, 'No NS record found'
        authSet = res.response.authority.pop()
        nsIPs = []
        if authSet:
            for ns in authSet:
                nsIp = res.response.get_rrset(message.ADDITIONAL, dns.name.from_text(ns.to_text()), RdataClass.IN,
                                              RdataType.A)
                if nsIp:
                    nsIPs.append(nsIp.pop())

        if not nsIPs:
            nsName = ns.to_text()
            print(f'Getting NS IP: {nsName}')
            res = dns.resolver.resolve(nsName)
            nsIp = res.response.get_rrset(message.ANSWER, dns.name.from_text(ns.to_text()), RdataClass.IN,
                                          RdataType.A).pop()
            print(f'nsIP: {nsIp}')
            nsIPs.append(nsIp)
            time.sleep(1)

        resolver.nameservers = [nsIPs[-1].to_text()]
        res = resolver.resolve(name)

    print(f'authority NS: {ns}')

    rrSet = res.response.get_rrset(message.ANSWER, name, RdataClass.IN, RdataType.A)
    if rrSet is not None:
        ip: dns.rdtypes.IN.A.A
        As = [ip.to_text() for ip in rrSet]
    else:
        As = []

    print(f'Real name: {name}')
    print(f'Real A: {As}')
    print(f'ttl: {rrSet.ttl}')
    return {
        "name": domain,
        "realName": name.to_text(),
        "A": As,
        "ttl": rrSet.ttl
    }


async def saveAuthorityTTLs(filename: str, limit=-1):
    with open(filename, 'r', encoding='utf-8') as f:
        lines = f.readlines()[:limit] if limit > 0 else f.readlines()
        coro = [saveOneDomain(line) for line in lines]
        await asyncio.gather(*coro)


def saveOneDomain(line):
    name, domain = line.rstrip().split(' ')
    res = getTTL(domain)
    res['cnName'] = name
    collection_domains.update_one({'name': res['name']}, {'$set': res}, upsert=True)


def getRRSet(domain: dns.name.Name, resolver: dns.resolver.Resolver):
    try:
        res = resolver.resolve(domain)
    except dns.resolver.LifetimeTimeout as e:
        # logger.debug(f'time out: {nameserver}')
        collection_timeout.update_one({
            'ip': resolver.nameservers[0]
        }, {
            '$set': {
                'available': False,
                'reason': str(e),
                'updateTime': timestampToTime(time.time())
            }
        }, upsert=True)
        return
    rrSet = res.response.get_rrset(message.ANSWER, domain, RdataClass.IN, RdataType.A)
    return rrSet


def probeQueryRate(domain: str, name: str, totalTTL: int, nameserver: str, comment: str, times: int = 5):
    global total_ip
    total_ip += 1
    domain = dns.name.from_text(domain)
    window = totalTTL / 10
    pre_tr = None
    interval = 0
    t_pre = 0
    net_delay = 0
    opt_time = 0
    try:
        for i in range(times):

            resolver = dns.resolver.Resolver(configure=False)
            resolver.nameservers = [nameserver]

            with sem:
                if t_pre:
                    opt_time = time.time() - t_pre
                    interval -= net_delay + opt_time
                    if interval < 0:
                        logger.error(f'[{comment}][#{i + 1}][{nameserver}]: {domain.to_text()}, '
                                     f'interval: {interval}s, net_delay: {net_delay}s, opt_time: {opt_time}s')
                    else:
                        time.sleep(interval)

                logger.debug(f'[{comment}][#{i + 1}][{nameserver}]: {domain.to_text()}, '
                             f'interval: {interval}s, net_delay: {net_delay}s, opt_time: {opt_time}s')

                query_time = time.time()
                rrSet = getRRSet(domain, resolver)
                net_delay = time.time() - query_time

                t_pre = time.time()

            if rrSet is None:
                return

            remain_ttl = rrSet.ttl
            tr = query_time - (totalTTL - remain_ttl)

            logger.debug(f'[{comment}][#{i + 1}]query time(Tp): {query_time}')
            logger.debug(f'[{comment}][#{i + 1}]remain ttl(Tl): {remain_ttl}')
            logger.debug(f'[{comment}][#{i + 1}]tr: {tr}')

            if remain_ttl == totalTTL:
                logger.info(f'[{comment}][#{i + 1}]Cache was refreshed by us')
                delta_tr = window
            else:
                delta_tr = (tr - pre_tr - totalTTL) if pre_tr else None

            interval = window + remain_ttl

            if interval > window + totalTTL:
                logger.error(f'[{comment}][#{i + 1}]interval: {interval} is larger than '
                             f'window + totalTTL: {window + totalTTL}, skip probing')
                return

            if delta_tr is not None and delta_tr < 0:
                logger.warning(f'[{comment}][#{i + 1}]negative delta_tr, sleep for interval: {interval}')

            else:
                logger.info(f'[{comment}][#{i + 1}]delta_tr: {delta_tr}')
                if delta_tr is not None:
                    ip: dns.rdtypes.IN.A.A
                    As = [ip.to_text() for ip in rrSet]
                    collection.insert_one({
                        'ip': nameserver,
                        'domain': domain.to_text(),
                        'cn_name': name,
                        'A': As,
                        'ttl': totalTTL,
                        'tr': tr,
                        'pre_tr': pre_tr,
                        'delta_tr': delta_tr,
                        'query_time': query_time,
                        'remain_ttl': remain_ttl,
                        'interval': interval,
                        'window': window,
                        'updateTime': timestampToTime(time.time())
                    })

            pre_tr = tr

    except Exception as e:
        logger.warning(f'[{comment}]nameserver: {nameserver} failed: {e}')
        try:
            collection_timeout.update_one({
                'ip': nameserver
            }, {
                '$set': {
                    'available': False,
                    'reason': str(e),
                    'updateTime': timestampToTime(time.time())
                }
            }, upsert=True)
        except:
            pass
        return


def main():
    records = collection_domains.find()
    with open('data/resolverIps.txt', 'r') as f:
        ips = list(map(str.rstrip, f.readlines()))
    ip_list = []
    step = 100
    for record in records:
        for i in range(0, len(ips), step):
            ip_list.append((record, ips[i:i + step]))

    with multiprocessing.Pool() as pool:
        pool.starmap(run, ip_list)


def run(record: dict, ips):
    realName = record['realName']
    totalTTL = record['ttl']
    if totalTTL > MAX_INTERVAL:
        return

    threads = []
    count = 1
    name = record['cnName']
    pid = multiprocessing.current_process().pid
    for i in range(len(ips)):
        logger.debug(f'#{pid} creating thread: #{count}')
        t = threading.Thread(target=probeQueryRate, args=(realName, name, totalTTL, ips[i], f'{pid}:{name}', 10))
        t.daemon = True
        t.start()
        threads.append(t)
        count += 1

    for t in threads:
        t.join()


if __name__ == '__main__':
    if sys.platform != "win32":
        os.system(f'ulimit -n {MAX_PORTS + 100}')
    main()
